﻿#include "RotatedShapeBasedMatching.h"
#include <QDir>
#include <QDebug>
#define _USE_MATH_DEFINES
#include "math.h"
namespace Qly
{

RotatedShapeBasedMatching::RotatedShapeBasedMatching()
{
}

bool RotatedShapeBasedMatching::trainModel(const cv::Mat &modelImg, std::vector<float> angles, int num_features)
{
    if(modelImg.empty())
    {
        return false;
    }

    preprocessModelImage(modelImg, angles);

    infos_have_templ.clear();
    shape_based_matching::shapeInfo_producer shapes;
    shapes.src = padded_templateImg;
    shapes.mask = padded_mask;
    shapes.produce_infos(angles, 1.0);

    if(pDetector)
    {
        delete pDetector;
    }
    pDetector = new line2Dup::Detector(num_features, {4, 8});
    bool is_first = true;
    int first_id = 0;
    float first_angle = 0;
    for (auto& info : shapes.infos) // 每个 info 对应一个旋转和缩放尺度
    {
        if (is_first)
        {
            info.templ_id = pDetector->addTemplate(shapes.src_of(info),
                                            class_id,
                                            shapes.mask_of(info),
                                            num_features); // 第一个当作是不旋转的模板
            first_id = info.templ_id;
            first_angle = info.angle;
            is_first = false;
            if(info.templ_id == -1)
            {
                std::cerr << "ShapeBasedMatching::trainModel cannot extract features from template image";
                return false;
            }
            buildFeatureImage(info);
        }
        else
        {
            info.templ_id = pDetector->addTemplate_rotate(class_id,
                                                   first_id,
                                                   info.angle - first_angle,
                                                   cv::Point2f{ shapes.src.cols / 2.0f, shapes.src.rows / 2.0f });
        }
        std::cout << "templ_id: " << info.templ_id << std::endl;
        if (info.templ_id != -1)
        {
            infos_have_templ.push_back(info);
        }
    }
    return true;
}

bool RotatedShapeBasedMatching::saveModel(std::string path, std::string modelName)
{
    if(pDetector == nullptr) return false;
    std::string fullModelName = path + "/" + modelName + ".yaml";
    std::string fullInfoName = path + "/" + modelName + "_info.yaml";
    std::string fullModelImageName = path + "/" + modelName + ".png";
    QDir dir;
    dir.remove(QString::fromStdString(fullModelName));
    dir.remove(QString::fromStdString(fullInfoName));

    // save templates
    pDetector->writeClasses(fullModelName);
    shape_based_matching::shapeInfo_producer::save_infos(infos_have_templ, fullInfoName);

    //std::cout << "save model end" << std::endl << std::endl;
    return cv::imwrite(fullModelImageName, templateImg);
}

bool RotatedShapeBasedMatching::loadModel(std::string path, std::string modelName)
{
    //模板文件路径获取
    std::string fullModelName = path + "/" + modelName + ".yaml";
    std::string fullInfoName = path + "/" + modelName + "_info.yaml";
    std::string fullModelImageName = path + "/" + modelName + ".png";
    class_ids.clear();
    class_ids.push_back("test");

    pDetector->readClasses(class_ids, fullModelName);
    infos_have_templ = shape_based_matching::shapeInfo_producer::load_infos(fullInfoName);

    std::vector<float> angles;
    for(auto it = infos_have_templ.begin(); it != infos_have_templ.end(); ++it)
    {
        angles.push_back((*it).angle);
    }
    templateImg = cv::imread(fullModelImageName);

    cv::Size2f sz = minSize(angles, templateImg.size());
    m_width_padding = (sz.width - templateImg.cols + 1) / 2;
    m_height_padding = (sz.height - templateImg.rows + 1) / 2;
    if(pDetector->numTemplates() != 0)
    {
        return true;
    }
    return false;
}

std::vector<RotatedShapeBasedMatching::MatchResult> RotatedShapeBasedMatching::runMatch(const cv::Mat &img, float score, size_t maxCount)
{
    std::vector<ShapeBasedMatching::MatchResult> ret;
    if(img.empty() || !pDetector) return ret;
    //resultImg = img.clone();
    cv::Mat mat = preprocessTestImage(img);
   // cv::imshow("", mat);
    std::vector<line2Dup::Match> matches = pDetector->match(mat, score, class_ids);// 这里得到的 matches 已经是排好序的了

    qDebug() <<"runMatch: pDetector->match size = " << matches.size();
    for(size_t i = 0; i < matches.size(); i++)
    {
        qDebug() << "runMatch: angle = " << -infos_have_templ[matches[i].template_id].angle
                 << ", score = " << matches[i].similarity << ", id = " << matches[i].template_id;
    }

    if(matches.size() == 0) return ret;
    refine(matches, maxCount, ret);
    return ret;
}

void RotatedShapeBasedMatching::preprocessModelImage(const cv::Mat &modelImage, std::vector<float> angles)
{
    cv::Size2f sz = minSize(angles, modelImage.size());
    m_width_padding = (sz.width - modelImage.cols + 1) / 2;
    m_height_padding = (sz.height - modelImage.rows + 1) / 2;

    ShapeBasedMatching::preprocessModelImage(modelImage, m_width_padding, m_height_padding);
}

cv::Size2f RotatedShapeBasedMatching::minSize(std::vector<float> angles, cv::Size imageSize)
{
    float width = imageSize.width;
    float height = imageSize.height;

    float x = width / 2.0;
    float y = height / 2.0;

    for(auto it = angles.begin(); it != angles.end(); ++it)
    {
        float theta = (*it) / 180.0 * M_PI;
        float cos_t = cos(theta);
        float sin_t = sin(theta);
        float new_width = 2 * fabs(x * cos_t + y * sin_t);
        float new_height = 2 * fabs(- y * sin_t + x * cos_t);
        if(width < new_width ) width = new_width;
        if(height < new_height) height = new_height;
    }
    return cv::Size2f(width, height);
}

}
